#!/usr/bin/env python3
import rclpy  # ROS 2 Python客户端库
from rclpy.node import Node  # ROS 2节点基类
from livox_interfaces.msg import CustomMsg  # Livox LiDAR自定义消息类型
from sensor_msgs.msg import PointCloud2, PointField  # ROS标准点云消息
import numpy as np  # 用于高效数值计算
import struct  # 用于二进制数据打包/解包
import open3d as o3d  # 用于3D点云处理和可视化
from threading import Thread, Lock  # 用于多线程编程
import time  # 用于时间测量和延时

# 定义PointCloud2的字段（X、Y、Z、反射率）
FIELDS = [
    PointField(name='x', offset=0, datatype=PointField.FLOAT32, count=1),
    PointField(name='y', offset=4, datatype=PointField.FLOAT32, count=1),
    PointField(name='z', offset=8, datatype=PointField.FLOAT32, count=1),
    PointField(name='intensity', offset=12, datatype=PointField.UINT8, count=1),
]
POINT_STEP = 13  # 每个点的字节数（4+4+4+1）
ROW_STEP = 0  # 每行字节数（动态计算）

class LidarProcessor(Node):
    def __init__(self):
        super().__init__('lidar_processor')  # 初始化ROS 2节点，名称为"lidar_processor"
        
        # 声明并获取ROS参数：输入点云话题名称
        self.declare_parameter('input_cloud', '/livox/lidar_avia')
        input_topic = self.get_parameter('input_cloud').get_parameter_value().string_value
        
        # 创建订阅者，订阅LiDAR数据，回调函数为self.callback
        self.sub = self.create_subscription(
            CustomMsg, input_topic, self.callback, 10
        )
        
        # 创建发布者，发布处理后的点云数据（ROS标准格式）
        self.pcl_pub = self.create_publisher(
            PointCloud2, '/livox/pointcloud2', 10
        )
        
        # 记录日志，显示订阅和发布的话题名称
        self.get_logger().info(f"订阅: {input_topic}，发布可视化点云到: /livox/pointcloud2")

        
        # 创建Open3D点云对象和可视化窗口
        self.pcd = o3d.geometry.PointCloud()
        self.vis = o3d.visualization.Visualizer()
        self.vis.create_window()
        
        # 设置可视化参数：点大小和背景颜色
        self.vis.get_render_option().point_size = 2.0
        self.vis.get_render_option().background_color = np.array([0, 0, 0])
        
        # 创建锁和标志变量，用于多线程同步
        self.point_cloud_lock = Lock()
        self.is_visualizing = True
        self.is_first_frame = True
        
        # 启动单独的线程用于点云可视化（避免阻塞主循环）
        self.visualization_thread = Thread(target=self.visualization_loop)
        self.visualization_thread.daemon = True  # 设置为守护线程
        self.visualization_thread.start()
        
        # 点云计数和保存相关参数
        self.frame_count = 0  # 帧计数器
        self.save_path = "/path/to/save/pointclouds/"  # 点云保存路径
        self.save_enabled = False  # 是否启用保存功能
        self.use_fixed_color = True  # 是否使用固定颜色（白色）显示点云
        
        # 过滤参数：距离阈值（50米）和角度阈值（30度）
        self.distance_threshold = 50.0  # 只保留50米内的点
        self.angle_threshold = 15.0     # 排除正前方30度范围内的点
        self.exclude_front = True      # 是否启用方向过滤
        
        # 预计算角度阈值的余弦值（用于后续计算）
        self.cos_angle_threshold = np.cos(np.radians(self.angle_threshold))

    def callback(self, msg: CustomMsg):
        """回调函数：处理接收到的点云消息"""
        start_time = time.time()  # 记录处理开始时间，用于性能分析
        
        # 将接收到的点云消息转换为NumPy数组，便于高效处理
        points_array = np.zeros((len(msg.points), 4), dtype=np.float32)
        for i, p in enumerate(msg.points):
            points_array[i, 0] = p.x  # X坐标
            points_array[i, 1] = p.y  # Y坐标
            points_array[i, 2] = p.z  # Z坐标
            points_array[i, 3] = p.reflectivity  # 反射率值
        
        # 执行距离过滤
        # 计算每个点到原点(0,0,0)的欧氏距离
        distances = np.sqrt(np.sum(points_array[:, :3]**2, axis=1))
        
        # 创建距离掩码：True表示距离小于等于阈值的点
        distance_mask = distances <= self.distance_threshold
        
        # 如果启用了方向过滤，则排除正前方指定角度范围内的点
        if self.exclude_front:
            # 定义正前方方向向量（X轴正方向）
            forward_vector = np.array([1.0, 0.0, 0.0])
            
            # 获取所有点的坐标向量
            point_vectors = points_array[:, :3]
            
            # 归一化点向量（将向量长度缩放到1）
            norms = np.linalg.norm(point_vectors, axis=1)
            norms[norms == 0] = 1  # 避免除以零
            normalized_vectors = point_vectors / norms[:, np.newaxis]
            
            # 计算每个点向量与正前方向量的点积（即夹角余弦值）
            cos_angles = np.dot(normalized_vectors, forward_vector)
            
            # 创建方向掩码：排除夹角余弦值绝对值小于阈值的点
            direction_mask = np.abs(cos_angles) < self.cos_angle_threshold
            
            # 合并距离和方向掩码：同时满足两个条件的点才保留
            filter_mask = np.logical_and(distance_mask, direction_mask)
        else:
            filter_mask = distance_mask  # 如果未启用方向过滤，仅使用距离掩码
        
        # 应用过滤掩码，保留符合条件的点
        filtered_points = points_array[filter_mask]
        
        # 如果过滤后没有有效点，记录日志并返回
        if len(filtered_points) == 0:
            self.get_logger().info(f"过滤后无有效点（距离阈值: {self.distance_threshold}m, 角度阈值: {self.angle_threshold}°）")
            return
        
        # 将过滤后的点云转换为ROS标准PointCloud2消息格式
        pcl_msg = self.create_pointcloud2_msg(msg.header, filtered_points)
        
        # 发布处理后的点云消息，供其他ROS节点使用（如RViz）
        self.pcl_pub.publish(pcl_msg)
        
        # 加锁以确保线程安全（防止在更新点云时进行渲染）
        with self.point_cloud_lock:
            points_np = filtered_points[:, :3]  # 提取坐标数据
            
            # 首次接收到点云时，检查坐标范围并记录日志
            if self.is_first_frame:
                x_min, y_min, z_min = np.min(points_np, axis=0)
                x_max, y_max, z_max = np.max(points_np, axis=0)
                self.get_logger().info(f"过滤后点云坐标范围: x=[{x_min:.2f}, {x_max:.2f}], y=[{y_min:.2f}, {y_max:.2f}], z=[{z_min:.2f}, {z_max:.2f}]")
                
                # 如果坐标值过大，自动应用缩放因子
                scale = 1.0
                if max(abs(x_max), abs(x_min), abs(y_max), abs(y_min), abs(z_max), abs(z_min)) > 1000:
                    scale = 0.001
                    self.get_logger().info(f"点云坐标值过大，将应用缩放因子: {scale}")
                    points_np *= scale
            
            # 更新Open3D点云对象的坐标
            self.pcd.points = o3d.utility.Vector3dVector(points_np)
            
            # 设置点云颜色
            if self.use_fixed_color:
                # 使用固定的白色
                colors = np.ones((len(filtered_points), 3))
            else:
                # 基于反射率设置颜色（反射率越高，红色越亮）
                intensities = filtered_points[:, 3] / 255.0  # 归一化反射率
                colors = np.zeros((len(filtered_points), 3))
                colors[:, 0] = intensities  # 红色通道表示反射率
            
            # 更新Open3D点云对象的颜色
            self.pcd.colors = o3d.utility.Vector3dVector(colors)
            
            # 如果启用保存功能，则保存当前帧点云
            if self.save_enabled:
                self.save_current_frame()
            
            # 标记不再是第一帧
            self.is_first_frame = False
        
        # 更新帧计数器
        self.frame_count += 1
        
        # 计算处理时间并记录日志
        process_time = (time.time() - start_time) * 1000  # 转换为毫秒
        self.get_logger().info(f"接收到点云: 原始{len(msg.points)}点，过滤后{len(filtered_points)}点，\
                               帧号: {self.frame_count}，处理时间: {process_time:.2f}ms")

    def create_pointcloud2_msg(self, header, points_array):
        """创建PointCloud2消息"""
        pcl_msg = PointCloud2()
        pcl_msg.header = header  # 复制原始消息的头部信息
        pcl_msg.height = 1  # 无序点云高度为1
        pcl_msg.width = len(points_array)  # 点云宽度为点的数量
        pcl_msg.fields = FIELDS  # 设置字段定义
        pcl_msg.is_bigendian = False  # 小端字节序
        pcl_msg.point_step = POINT_STEP  # 每个点的字节数
        pcl_msg.row_step = POINT_STEP * pcl_msg.width  # 每行字节数
        pcl_msg.is_dense = False  # 可能包含无效点（NaN或Inf）
        
        # 将点数据打包为二进制字节流
        data = []
        for point in points_array:
            x, y, z, intensity = point
            data.append(struct.pack('fffB', x, y, z, int(intensity)))
        pcl_msg.data = b''.join(data)  # 合并所有点的二进制数据
        
        return pcl_msg

    def visualization_loop(self):
        """可视化线程：持续更新点云显示"""
        is_first_visualization = True
        
        while self.is_visualizing:  # 只要可视化标志为True就继续循环
            # 加锁以确保线程安全
            with self.point_cloud_lock:
                if is_first_visualization and len(self.pcd.points) > 0:
                    # 首次可视化且有有效点时，添加点云并设置初始视图
                    self.vis.add_geometry(self.pcd)
                    
                    # 设置视图参数：从Z轴上方俯视
                    ctr = self.vis.get_view_control()
                    ctr.set_front([0, 0, -1])  # 相机前向向量（从屏幕外向屏幕内）
                    ctr.set_up([0, 1, 0])      # 相机上方向向量（Y轴正方向）
                    ctr.set_lookat([0, 0, 0])  # 观察点为原点
                    ctr.set_zoom(0.8)         # 设置缩放比例
                    
                    is_first_visualization = False
                elif not is_first_visualization:
                    # 非首次可视化，更新点云并刷新视图
                    self.vis.update_geometry(self.pcd)
                    self.vis.poll_events()
                    self.vis.update_renderer()
            
            # 休眠一小段时间，控制刷新频率（约30FPS）
            time.sleep(0.03)

    def save_current_frame(self):
        """保存当前帧点云到文件"""
        filename = f"{self.save_path}/pointcloud_{self.frame_count:06d}.pcd"
        o3d.io.write_point_cloud(filename, self.pcd)
        self.get_logger().info(f"已保存点云: {filename}")

    def destroy_node(self):
        """节点销毁时的清理工作"""
        self.is_visualizing = False  # 停止可视化循环
        self.vis.destroy_window()    # 关闭可视化窗口
        super().destroy_node()       # 调用父类的销毁方法

def main(args=None):
    rclpy.init(args=args)  # 初始化ROS 2 Python客户端库
    node = LidarProcessor()  # 创建节点实例
    
    try:
        rclpy.spin(node)  # 进入节点的主循环，等待消息
    except KeyboardInterrupt:
        pass  # 捕获Ctrl+C中断信号
    finally:
        node.destroy_node()  # 销毁节点，执行清理工作
        rclpy.shutdown()  # 关闭ROS 2客户端库

if __name__ == '__main__':
    main()  # 程序入口点

